# -*- coding: utf-8 -*-

from agent import Agent
from simulator import Simulator

from pettingzoo.magent import battle_v2

max_epoch = 10
max_episode = 100
max_step = 100


def run(env):
    """ Training loop """

    """ Step 1: Initialize agent and roll-out simulator. """
    agent = Agent(agent_ls=env.possible_agents,
                  n_obs=env.observation_spaces,
                  n_action=list(env.action_spaces.values())[0].n)
    simulator = Simulator(env=env, agent=agent)

    for _ in range(max_epoch):  # Train n times

        """ Step 1: Generate replay buffer. """
        replay_buffer = simulator.generate_replays(num_episode=max_episode, num_step=max_step)

        """ Step 2: Update agent networks. """
        agent.train(replay_buffer)


if __name__ == '__main__':

    test_env = battle_v2.parallel_env()

    run(env=test_env)

